import logging
import time
from datetime import datetime
from tqdm import tqdm
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from pinecone import Pinecone, ServerlessSpec
from pinecone.exceptions import PineconeApiException  # 导入异常类
from collections import Counter

# 配置logging，包含日期和时间
logging.basicConfig(
    format='%(asctime)s - %(levelname)s - %(message)s',
    level=logging.INFO,
    datefmt='%Y-%m-%d %H:%M:%S'
)

# 初始化Pinecone客户端（使用新API密钥）
pc = Pinecone(api_key="pcsk_Njyyj_Jx3sovkra1ygNdaXR6vXicgz8s5fiWnS6dbCxTjjr1bwqvZtYeBXxobphF2TiX8")
index_name = "mnist-index"

# 检查索引是否存在（带重试机制，解决同步延迟）
def index_exists(index_name):
    for _ in range(3):  # 最多重试3次
        try:
            return index_name in pc.list_indexes()
        except Exception as e:
            logging.warning(f"检查索引存在性失败，重试... 错误: {str(e)}")
            time.sleep(1)  # 等待1秒后重试
    return False

# 确保索引存在且可连接
if not index_exists(index_name):
    logging.info(f"索引 '{index_name}' 不存在，尝试创建...")
    try:
        pc.create_index(
            name=index_name,
            dimension=64,
            metric="euclidean",
            spec=ServerlessSpec(
                cloud="aws",
                region="us-east-1"
            )
        )
        # 等待索引创建完成（Serverless索引需要几秒初始化）
        time.sleep(5)
        logging.info(f"索引 '{index_name}' 创建成功")
    except PineconeApiException as e:
        if "ALREADY_EXISTS" in str(e):
            logging.info(f"索引 '{index_name}' 已存在，直接连接")
        else:
            raise e  # 其他错误则抛出
else:
    logging.info(f"索引 '{index_name}' 已存在，直接连接")

# 连接到索引
index = pc.Index(index_name)

# 加载数据集并拆分
digits = load_digits(n_class=10)
X = digits.data
y = digits.target
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

# 上传训练数据（仅当索引为空时）
try:
    index_stats = index.describe_index_stats()
    if index_stats['total_vector_count'] == 0:
        logging.info("索引为空，开始上传训练数据...")
        batch_size = 100
        total_vectors = len(X_train)
        # 构造向量列表
        vectors = [
            (f"train_{i}", X_train[i].tolist(), {"label": int(y_train[i])})
            for i in range(total_vectors)
        ]
        # 带进度条批量上传
        for i in tqdm(range(0, total_vectors, batch_size), desc="上传进度"):
            index.upsert(vectors[i:i + batch_size])
        logging.info(f"成功上传{total_vectors}条数据")
    else:
        logging.info(f"索引已包含{index_stats['total_vector_count']}条数据，跳过上传")
except Exception as e:
    logging.error(f"获取索引状态失败: {str(e)}")
    raise

# 测试准确率（k=11）
correct = 0
total = len(X_test)
for i in tqdm(range(total), desc="测试进度"):
    # 准备查询向量
    query_vector = X_test[i].tolist()
    # 检索最相似的11个向量
    results = index.query(
        vector=query_vector,
        top_k=11,
        include_metadata=True
    )
    # 投票机制获取预测结果
    labels = [match['metadata']['label'] for match in results['matches']]
    prediction = Counter(labels).most_common(1)[0][0]
    # 统计正确数
    if prediction == y_test[i]:
        correct += 1

# 计算并输出准确率
accuracy = correct / total * 100
logging.info(f"当k=11时，使用Pinecone的准确率: {accuracy:.2f}%")